from sqlalchemy import Column, Integer, String, DateTime, func
from datetime import datetime
from api_server.db.base import Base
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    Literal,
    Optional,
    Callable,
    Generator,
    Dict,
    Any,
    Awaitable,
    Union,
    Tuple
)


class KnowledgeBaseModel(Base):
    """
    知识库模型
    """
    __tablename__ = 'knowledge_base'
    id = Column(Integer, primary_key=True, autoincrement=True, comment='知识库ID')
    kb_name = Column(String(50), comment='知识库名称,只能为英文')
    kb_nickname = Column(String(50), default='', comment='知识库昵称可以为中文')
    kb_info = Column(String(200), comment='知识库简介(用于Agent)')
    vs_type = Column(String(50), comment='向量库类型')
    kb_type = Column(String(50), default='public', comment='知识库类型(public(公共),person(个人))')
    kb_prop = Column(String(50), default='comm', comment='知识库属性(sys(系统的) 属于公共范围,comm普通)')
    embed_model = Column(String(50), comment='嵌入模型名称')
    file_count = Column(Integer, default=0, comment='文件数量')
    create_user_id = Column(Integer, default=0, comment='创建者id')
    create_user_name = Column(String(50), default='', comment='创建者名称(冗余)')
    create_time = Column(DateTime, default=datetime.now(), comment='创建时间')

    def __repr__(self):
        return f"<KnowledgeBase(id='{self.id}', kb_name='{self.kb_name}',kb_intro='{self.kb_info} vs_type='{self.vs_type}', embed_model='{self.embed_model}', file_count='{self.file_count}', create_time='{self.create_time}')>"


from typing import List
from api_server.utils import BaseResponse
import pydantic
from pydantic import BaseModel, BaseSettings
from api_server.auth_util import settings


class KnowledgeBean(BaseModel):
    id: int = pydantic.Field(1, description='知识库ID')
    kb_name: str = pydantic.Field('sample', description='知识库名称')
    kb_nickname: Any = pydantic.Field('我是规章制度知识库', description='知识库昵称')
    kb_info: Any = pydantic.Field('这是一个很好的知识库', description='知识库简介(用于Agent)')
    kb_prop: Any = pydantic.Field('sys', description='知识库属性(sys(系统的) 属于公共范围)')
    vs_type: str = pydantic.Field('pg', description='向量库类型')
    kb_type: Any = pydantic.Field('public', description='知识库类型(public(公共),person(个人),sys(系统的))')
    embed_model: str = pydantic.Field('bge-large-zh-v1.5', description='嵌入模型名称')
    file_count: int = pydantic.Field(None, description='文件数量')
    create_user_id: Any = pydantic.Field(1, description='创建者id')
    create_user_name: Any = pydantic.Field("admin", description='创建者名称(冗余)')
    create_time: Any = pydantic.Field(None, description='创建时间')

    class Config:
        orm_mode = True
        json_encoders = {
            datetime: lambda v: v.strftime(settings.datetime_format)
        }


class ListKbResponse(BaseResponse):
    data: List[KnowledgeBean] = pydantic.Field(..., description="知识库对象")

    class Config:
        schema_extra = {
            "example": {
                "code": 200,
                "msg": "success",
                "data": [{
                    "id": 2,
                    "kb_name": "temphetong",
                    "kb_info": "关于temphetong的知识库",
                    "vs_type": "pg",
                    "kb_type": "null",
                    "embed_model": "qwen-api",
                    "file_count": 4,
                    "create_user_id": "null",
                    "create_user_name": "null",
                    "create_time": "2024-05-14T07:27:48.191640"
                },
                    {
                        "id": 3,
                        "kb_name": "chanwd",
                        "kb_info": "关于chanwd的知识库",
                        "vs_type": "pg",
                        "kb_type": "null",
                        "embed_model": "qwen-api",
                        "file_count": 0,
                        "create_user_id": "null",
                        "create_user_name": "null",
                        "create_time": "2024-05-15T10:08:04.523757"
                    },
                    {
                        "id": 4,
                        "kb_name": "test232",
                        "kb_info": "关于test232的知识库",
                        "vs_type": "pg",
                        "kb_type": "null",
                        "embed_model": "bge-large-zh-v1.5",
                        "file_count": 0,
                        "create_user_id": 0,
                        "create_user_name": "0",
                        "create_time": "2024-05-20T08:00:19.985560"
                    },
                    {
                        "id": 5,
                        "kb_name": "test24232",
                        "kb_info": "关于test24232的知识库",
                        "vs_type": "pg",
                        "kb_type": "null",
                        "embed_model": "qwen-api",
                        "file_count": 0,
                        "create_user_id": 0,
                        "create_user_name": "0",
                        "create_time": "2024-05-20 08:20:23.408469"
                    }],
            }
        }


if __name__ == "__main__":
    data = {
        "id": 5,
        "kb_name": "test24232",
        "kb_info": "关于test24232的知识库",
        "vs_type": "pg",
        "kb_type": "null",
        "embed_model": "qwen-api",
        "file_count": 0,
        "create_user_id": 0,
        "create_user_name": "0",
        "create_time": "2024-05-20 08:20:23"
    }
    # temp = KnowledgeBean(**data)
    temp = KnowledgeBaseModel(**data)
    print(temp.__dict__)
